import cv2
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt


#定义一个对检测结果的判定，如果检测错误，报错并停止执行
def check_result(num):
    try:
        # 使用assert语句检查条件是否为真
        assert num == 2, "检测错误：环境干扰较多"
    except AssertionError as e:
        # 捕获AssertionError异常并处理
        print(f"错误信息：{e}")
        # 抛出异常，终止程序
        raise

path = 'Hand0008.jpg'

img = cv2.imread(path)
#获取原始图像的宽高
original_height,original_width = img.shape[:2]
#计算比例
scale = min(1920/original_width,1280/original_height)

#计算新宽高
new_width = int(original_width * scale)
new_height = int(original_height * scale)

#resize原始图像，复原宽高比
img = cv2.resize(img,(new_width,new_height))

img_HSV = cv2.cvtColor(img,cv2.COLOR_BGR2HSV)

#定义红色范围
lower_red = np.array([0,80,80])
upper_red = np.array([10,255,255])
#创建红色蒙版
red_mask0 = cv2.inRange(img_HSV,lower_red,upper_red)

lower_red = np.array([160,80,80])
upper_red = np.array([180,255,255])
#创建红色蒙版
red_mask1 = cv2.inRange(img_HSV,lower_red,upper_red)

#合并红色蒙版
red_mask = red_mask0 + red_mask1

kernel = np.ones((10,10),np.uint8)
red_mask_dilated = cv2.dilate(red_mask,kernel,iterations=4)

# 找膨胀之后红线的外框
contours, _ = cv2.findContours(red_mask_dilated,cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_SIMPLE)
num_of_line = len(contours)
#判断红色线条是否提取正确
check_result(num_of_line)



#找红线的外框外框的中心点
screw_centers = []
for contour in contours:
        x, y, w, h = cv2.boundingRect(contour)

        cv2.rectangle(img,(x,y),(x+w,y+h),(0,255,0),2)

        center_x = x + w // 2
        center_y = y + h // 2
        screw_centers.append((center_x,center_y))

        # 截取红线外框ROI,像素下标会从0开始，后续计算需要做补偿
        roi = red_mask[y:y + h, x:x + w]  #y:x 第一个是纵向，第二个是横向



        # 查找二值图像中白色像素的坐标
        white_pixels = np.column_stack(np.where(roi > 0))

        white_pixels[:, 0] = white_pixels[:, 0] + y
        white_pixels[:, 1] = white_pixels[:, 1] + x

        print(white_pixels)
        array_shape = white_pixels.shape
        print("数组维度:", array_shape)

        # 进行直线拟合，最少需要两个点
        if len(white_pixels) >= 2:
                # 一次多项式拟合，直线
                line_params = np.polyfit(white_pixels[:, 0], white_pixels[:, 1], 1)
                print("line_params:", line_params)
                # 带入x坐标范围，获取拟合的直线方程算出来的值
                line_y = np.polyval(line_params, np.arange(x, x + w))

                print(f"Line coordinates: x1={x}, y1={int(line_y[0])}, x2={x + w - 1}, y2={int(line_y[-1])}")

                # cv2.line(img,(x, int(line_y[1])),(x+w , int(line_y[w-1]) ),(255,0,0),2 )
                for i in range(len(line_y) - 1):
                        cv2.line(img, (x + i, int(line_y[i])), (x + i + 1, int(line_y[i + 1])), (255, 0, 0), 2)

        cv2.circle(img, (center_x, center_y), radius=5, color=(255, 0, 0), thickness=-1)  # -1表示填充圆
        print(f"红线外接框中心坐标为:{center_x},{center_y}")



#提取白色区域像素，进行拟合

cv2.namedWindow('img',cv2.WINDOW_NORMAL) #解决显示过大，显示不全的问题，解决图片分辨率大于电脑分辨率问题
cv2.imshow('img',img)
cv2.imshow('red_mask',red_mask)
cv2.waitKey(0)
cv2.destroyWindow('img')
cv2.destroyWindow('red_mask')


# white_pixels[:, 0] = white_pixels[:, 0] + x
# white_pixels[:, 1] = white_pixels[:, 1] + y

# ####################验证一下找到的区域在哪里######################################
#         # 将数组转换为 pandas DataFrame
#         df = pd.DataFrame(white_pixels, columns=['列1', '列2'])  # 列名可以根据实际情况更改
#
#         # 将 DataFrame 存储为 Excel 文件
#         excel_filename = 'output.xlsx'  # Excel 文件的名称
#         df.to_excel(excel_filename, index=False)
# ##############################################################
#
#
#
# ####################验证一下找到的区域在哪里######################################
#         # 判断白色区域是否有点
#         if len(white_pixels) > 0:
#             # 获取白色区域的边缘坐标
#             contours, _ = cv2.findContours(roi, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
#             # 填充ROI为红色
#             cv2.fillPoly(img, contours, color=(0, 0, 255))
# ##############################################################

#在两个提取出来的白色区域，拟合一条该直线的方程式，利用两个拟合的方程分别计算角度
# #1.霍夫变换
# lines = cv2.HoughLines(red_mask_dilated, 1, np.pi / 180, threshold=200)
# # 绘制检测到的直线
# for line in lines:
#     rho, theta = line[0]
#     a = np.cos(theta)
#     b = np.sin(theta)
#     x0 = a * rho
#     y0 = b * rho
#     x1 = int(x0 + 1000 * (-b))
#     y1 = int(y0 + 1000 * (a))
#     x2 = int(x0 - 1000 * (-b))
#     y2 = int(y0 - 1000 * (a))
#     cv2.line(img, (x1, y1), (x2, y2), (0, 0, 255), 2)